#!/usr/bin/env bash

set -eu -o pipefail
shopt -qs failglob

for opt in "$@"; do
   optarg="$(expr "${opt}" : '[^=]*=\(.*\)')"
   case "${opt}" in
      --package-dir=*) PACKAGE_DIR="${optarg}" ;;
      --output-dir=*) OUTPUT_DIR="${optarg}" ;;
   esac
done

mkdir -p "${OUTPUT_DIR}"

FILENAME_PREFIX="bottlerocket-${VARIANT}-${ARCH}-${VERSION_ID}-${BUILD_ID}"
DISK_IMAGE_NAME="${FILENAME_PREFIX}.img.lz4"
BOOT_IMAGE_NAME="${FILENAME_PREFIX}-boot.ext4.lz4"
VERITY_IMAGE_NAME="${FILENAME_PREFIX}-root.verity.lz4"
ROOT_IMAGE_NAME="${FILENAME_PREFIX}-root.ext4.lz4"
DATA_IMAGE_NAME="${FILENAME_PREFIX}-data.img.lz4"

DISK_IMAGE="$(mktemp)"
BOOT_IMAGE="$(mktemp)"
VERITY_IMAGE="$(mktemp)"
ROOT_IMAGE="$(mktemp)"
DATA_IMAGE="$(mktemp)"
EFI_IMAGE="$(mktemp)"
PRIVATE_IMAGE="$(mktemp)"
BOTTLEROCKET_DATA="$(mktemp)"

ROOT_MOUNT="$(mktemp -d)"
BOOT_MOUNT="$(mktemp -d)"
DATA_MOUNT="$(mktemp -d)"
EFI_MOUNT="$(mktemp -d)"

SELINUX_ROOT="/etc/selinux"
SELINUX_POLICY="fortified"
SELINUX_FILE_CONTEXTS="${ROOT_MOUNT}/${SELINUX_ROOT}/${SELINUX_POLICY}/contexts/files/file_contexts"

VERITY_VERSION=1
VERITY_HASH_ALGORITHM=sha256
VERITY_DATA_BLOCK_SIZE=4096
VERITY_HASH_BLOCK_SIZE=4096

# Define partition type GUIDs for all OS-managed partitions. This is required
# for the boot partition, where we set gptprio bits in the GUID-specific use
# field, but we might as well do it for all of them.
BOTTLEROCKET_BOOT_TYPECODE="6b636168-7420-6568-2070-6c616e657421"
EFI_SYSTEM_TYPECODE="C12A7328-F81F-11D2-BA4B-00A0C93EC93B"
BOTTLEROCKET_ROOT_TYPECODE="5526016a-1a97-4ea4-b39a-b7c8c6ca4502"
BOTTLEROCKET_HASH_TYPECODE="598f10af-c955-4456-6a99-7720068a6cea"
BOTTLEROCKET_RESERVED_TYPECODE="0c5d99a5-d331-4147-baef-08e2b855bdc9"
BOTTLEROCKET_PRIVATE_TYPECODE="440408bb-eb0b-4328-a6e5-a29038fad706"
BOTTLEROCKET_DATA_TYPECODE="626f7474-6c65-6474-6861-726d61726b73"

if [[ "${ARCH}" == "x86_64" ]]; then
  FIRM_NAME="BIOS-BOOT"
  BOTTLEROCKET_FIRM_TYPECODE="ef02"
else
  FIRM_NAME="EFI_BOOT"
  BOTTLEROCKET_FIRM_TYPECODE="${EFI_SYSTEM_TYPECODE}"
fi

truncate -s 2G "${DISK_IMAGE}"
# boot: 40M + root: 920M + hash: 10M + reserved: 30M = 1000M
# boot partition attributes (-A): 48 = gptprio priority bit; 56 = gptprio successful bit
# partitions are backwards so that we don't make things inconsistent when specifying a wrong end sector :)
sgdisk --clear \
   -n 0:2005M:2047M -c 0:"BOTTLEROCKET-PRIVATE"    -t 0:"${BOTTLEROCKET_PRIVATE_TYPECODE}" \
   -n 0:1975M:0     -c 0:"BOTTLEROCKET-RESERVED-B" -t 0:"${BOTTLEROCKET_RESERVED_TYPECODE}" \
   -n 0:1965M:0     -c 0:"BOTTLEROCKET-HASH-B"     -t 0:"${BOTTLEROCKET_HASH_TYPECODE}" \
   -n 0:1045M:0     -c 0:"BOTTLEROCKET-ROOT-B"     -t 0:"${BOTTLEROCKET_ROOT_TYPECODE}" \
   -n 0:1005M:0     -c 0:"BOTTLEROCKET-BOOT-B"     -t 0:"${BOTTLEROCKET_BOOT_TYPECODE}" -A 0:"clear":48 -A 0:"clear":56 \
   -n 0:975M:0      -c 0:"BOTTLEROCKET-RESERVED-A" -t 0:"${BOTTLEROCKET_RESERVED_TYPECODE}" \
   -n 0:965M:0      -c 0:"BOTTLEROCKET-HASH-A"     -t 0:"${BOTTLEROCKET_HASH_TYPECODE}" \
   -n 0:45M:0       -c 0:"BOTTLEROCKET-ROOT-A"     -t 0:"${BOTTLEROCKET_ROOT_TYPECODE}" \
   -n 0:5M:0        -c 0:"BOTTLEROCKET-BOOT-A"     -t 0:"${BOTTLEROCKET_BOOT_TYPECODE}" -A 0:"set":48 -A 0:"set":56 \
   -n 0:1M:0        -c 0:"${FIRM_NAME}"    -t 0:"${BOTTLEROCKET_FIRM_TYPECODE}" \
   --sort --print "${DISK_IMAGE}"

rpm -iv --root "${ROOT_MOUNT}" "${PACKAGE_DIR}"/*.rpm
install -p -m 0644 /host/{COPYRIGHT,LICENSE-APACHE,LICENSE-MIT} "${ROOT_MOUNT}"/usr/share/licenses/
mksquashfs \
  "${ROOT_MOUNT}"/usr/share/licenses \
  "${ROOT_MOUNT}"/usr/share/bottlerocket/licenses.squashfs \
  -no-exports -all-root -comp zstd
rm -rf "${ROOT_MOUNT}"/var/lib "${ROOT_MOUNT}"/usr/share/licenses/*

if [[ "${ARCH}" == "x86_64" ]]; then
  SYS_ROOT="x86_64-bottlerocket-linux-gnu/sys-root"
  # MBR and BIOS-BOOT
  echo "(hd0) ${DISK_IMAGE}" > ${ROOT_MOUNT}/boot/grub/device.map
  "${ROOT_MOUNT}/sbin/grub-bios-setup" \
     --directory="${ROOT_MOUNT}/boot/grub" \
     --device-map="${ROOT_MOUNT}/boot/grub/device.map" \
     --root="hd0" \
     --skip-fs-probe \
     ${DISK_IMAGE}

  rm -vf "${ROOT_MOUNT}"/boot/grub/* "${ROOT_MOUNT}"/sbin/grub*
else
  SYS_ROOT="aarch64-bottlerocket-linux-gnu/sys-root"
  # For aarch64 we need an EFI partition instead, formatted
  # FAT32 with the .efi binary at the correct path, eg /efi/boot.
  # grub-mkimage has put bootaa64.efi at /boot/efi/EFI/BOOT
  mv ${ROOT_MOUNT}/boot/efi/* "${EFI_MOUNT}"

  # The 'recommended' size for the EFI partition is 100MB but our aarch64.efi
  # only takes up around 700KB, so this will suffice for now.
  dd if=/dev/zero of=${EFI_IMAGE} bs=1M count=4
  mkfs.vfat -I -S 512 "${EFI_IMAGE}" $((4*2048))
  mmd -i ${EFI_IMAGE} ::/EFI
  mmd -i ${EFI_IMAGE} ::/EFI/BOOT
  mcopy -i ${EFI_IMAGE} ${EFI_MOUNT}/EFI/BOOT/bootaa64.efi ::/EFI/BOOT
  dd if="${EFI_IMAGE}" of="${DISK_IMAGE}" conv=notrunc bs=1M seek=1

  # Create the grub directory which grub-bios-setup would have otherwise done.
  mkdir -p ${ROOT_MOUNT}/boot/grub
fi

# Now that we're done messing with /, move /boot out of it
mv "${ROOT_MOUNT}/boot"/* "${BOOT_MOUNT}"

# Set the Bottlerocket variant, version, and build-id
echo "PRETTY_NAME=\"Bottlerocket OS ${VERSION_ID}\"" >> ${ROOT_MOUNT}/${SYS_ROOT}/usr/lib/os-release
echo "VARIANT_ID=${VARIANT}" >> ${ROOT_MOUNT}/${SYS_ROOT}/usr/lib/os-release
echo "VERSION_ID=${VERSION_ID}" >> ${ROOT_MOUNT}/${SYS_ROOT}/usr/lib/os-release
echo "BUILD_ID=${BUILD_ID}" >> ${ROOT_MOUNT}/${SYS_ROOT}/usr/lib/os-release

# BOTTLEROCKET-ROOT-A
mkdir -p "${ROOT_MOUNT}/lost+found"
ROOT_LABELS=$(setfiles -n -d -F -m -r "${ROOT_MOUNT}" \
    "${SELINUX_FILE_CONTEXTS}" "${ROOT_MOUNT}" \
    | awk -v root="${ROOT_MOUNT}" '{gsub(root"/","/"); gsub(root,"/"); print "ea_set", $1, "security.selinux", $4}')
mkfs.ext4 -O ^has_journal -b "${VERITY_DATA_BLOCK_SIZE}" -d "${ROOT_MOUNT}" "${ROOT_IMAGE}" 920M
echo "${ROOT_LABELS}" | debugfs -w -f - "${ROOT_IMAGE}"
resize2fs -M "${ROOT_IMAGE}"
dd if="${ROOT_IMAGE}" of="${DISK_IMAGE}" conv=notrunc bs=1M seek=45

# BOTTLEROCKET-VERITY-A
truncate -s 8M "${VERITY_IMAGE}"
veritysetup_output="$(veritysetup format \
    --format "$VERITY_VERSION" \
    --hash "$VERITY_HASH_ALGORITHM" \
    --data-block-size "$VERITY_DATA_BLOCK_SIZE" \
    --hash-block-size "$VERITY_HASH_BLOCK_SIZE" \
    "${ROOT_IMAGE}" "${VERITY_IMAGE}" \
    | tee /dev/stderr)"
if ! stat -c %s "${VERITY_IMAGE}" | grep -q '^8388608$'; then
    echo "verity partition is larger than expected (8M)"
    exit 1
fi
VERITY_DATA_4K_BLOCKS="$(grep '^Data blocks:' <<<$veritysetup_output | awk '{ print $NF }')"
VERITY_DATA_512B_BLOCKS="$(($VERITY_DATA_4K_BLOCKS * 8))"
VERITY_ROOT_HASH="$(grep '^Root hash:' <<<$veritysetup_output | awk '{ print $NF }')"
VERITY_SALT="$(grep '^Salt:' <<<$veritysetup_output | awk '{ print $NF }')"
veritysetup verify "${ROOT_IMAGE}" "${VERITY_IMAGE}" "${VERITY_ROOT_HASH}"
dd if="${VERITY_IMAGE}" of="${DISK_IMAGE}" conv=notrunc bs=1M seek=965

# write GRUB config
cat <<EOF > ${BOOT_MOUNT}/grub/grub.cfg
set default="0"
set timeout="0"

menuentry "Bottlerocket OS ${VERSION_ID}" {
   linux (\$root)/vmlinuz root=/dev/dm-0 rootwait ro \\
       console=tty0 console=ttyS0 random.trust_cpu=on selinux=1 enforcing=1 \\
       systemd.log_target=journal-or-kmsg systemd.log_color=0 net.ifnames=0 \\
       biosdevname=0 dm_verity.max_bios=-1 dm_verity.dev_wait=1 \\
       dm-mod.create="root,,,ro,0 $VERITY_DATA_512B_BLOCKS verity $VERITY_VERSION PARTUUID=\$boot_uuid/PARTNROFF=1 PARTUUID=\$boot_uuid/PARTNROFF=2 \\
       $VERITY_DATA_BLOCK_SIZE $VERITY_HASH_BLOCK_SIZE $VERITY_DATA_4K_BLOCKS 1 $VERITY_HASH_ALGORITHM $VERITY_ROOT_HASH $VERITY_SALT 1 restart_on_corruption"
}
EOF

# BOTTLEROCKET-BOOT-A
mkfs.ext4 -O ^has_journal -d "${BOOT_MOUNT}" "${BOOT_IMAGE}" 40M
resize2fs -M "${BOOT_IMAGE}"
dd if="${BOOT_IMAGE}" of="${DISK_IMAGE}" conv=notrunc bs=1M seek=5

# BOTTLEROCKET-PRIVATE

# Targeted toward the current API server implementation.
# Relative to the ext4 defaults, we:
# - adjust the inode ratio since we expect lots of small files
# - retain the inode size to allow most settings to be stored inline
# - retain the block size to handle worse-case alignment for hardware
mkfs.ext4 -b 4096 -i 4096 -I 256 "${PRIVATE_IMAGE}" 42M
dd if="${PRIVATE_IMAGE}" of="${DISK_IMAGE}" conv=notrunc bs=1M seek=2005

# BOTTLEROCKET-DATA
truncate -s 1G "${DATA_IMAGE}"
sgdisk --clear \
   -n 0:1M:1023M -c 0:"BOTTLEROCKET-DATA" -t 0:"${BOTTLEROCKET_DATA_TYPECODE}" \
   --sort --print "${DATA_IMAGE}"
# If we build on a host with SELinux enabled, we could end up with labels that
# do not match our policy. Since we allow replacing the data volume at runtime,
# we can't count on these labels being correct in any case, and it's better to
# remove them all.
UNLABELED=$(find "${DATA_MOUNT}" \
    | awk -v root="${DATA_MOUNT}" '{gsub(root"/","/"); gsub(root,"/"); print "ea_rm", $1, "security.selinux"}')
mkfs.ext4 -d "${DATA_MOUNT}" "${BOTTLEROCKET_DATA}" 1022M
echo "${UNLABELED}" | debugfs -w -f - "${BOTTLEROCKET_DATA}"
dd if="${BOTTLEROCKET_DATA}" of="${DATA_IMAGE}" conv=notrunc bs=1M seek=1

sgdisk -v "${DISK_IMAGE}"
sgdisk -v "${DATA_IMAGE}"

lz4 -vc "${DISK_IMAGE}" >"${OUTPUT_DIR}/${DISK_IMAGE_NAME}"
lz4 -vc "${DATA_IMAGE}" >"${OUTPUT_DIR}/${DATA_IMAGE_NAME}"
lz4 -9vc "${BOOT_IMAGE}" >"${OUTPUT_DIR}/${BOOT_IMAGE_NAME}"
lz4 -9vc "${VERITY_IMAGE}" >"${OUTPUT_DIR}/${VERITY_IMAGE_NAME}"
lz4 -9vc "${ROOT_IMAGE}" >"${OUTPUT_DIR}/${ROOT_IMAGE_NAME}"
chown 1000:1000 "${OUTPUT_DIR}/${DISK_IMAGE_NAME}" \
    "${OUTPUT_DIR}/${DATA_IMAGE_NAME}" \
    "${OUTPUT_DIR}/${BOOT_IMAGE_NAME}" \
    "${OUTPUT_DIR}/${VERITY_IMAGE_NAME}" \
    "${OUTPUT_DIR}/${ROOT_IMAGE_NAME}"

# Clean up temporary files to reduce size of layer.
rm -f "${PACKAGE_DIR}"/*.rpm
rm -rf /tmp/*
